
from langchain_community.vectorstores import Chroma
from langchain_openai.embeddings import AzureOpenAIEmbeddings  # 导入嵌入模型
from langchain_community.document_loaders import PyPDFLoader
from tool import get_azure_endpoint, get_api_key, get_api_version
if __name__ == '__main__':
    embedding = AzureOpenAIEmbeddings(
        azure_endpoint=get_azure_endpoint().rstrip('/'),  # 移除尾部斜杠，只保留基础URL
        model="text-embedding-3-small",  # 重命名为 azure_deployment
        api_key=get_api_key(),
        api_version=get_api_version()
    )
    vectordb_chinese = Chroma(
        persist_directory="./docs/chroma/matplotlib",
        embedding_function=embedding,
    )

    question_chinese = "Matplotlib是什么?"
    docs_ss_chinese = vectordb_chinese.similarity_search(question_chinese, k=3)
    print("docs[0]: ")
    print(docs_ss_chinese[0].page_content[:100])
    print()
    print("docs[1]: ")
    print(docs_ss_chinese[1].page_content[:100])

    docs_mmr_chinese = vectordb_chinese.max_marginal_relevance_search(question_chinese, k=3)
    print("docs[0]: ")
    print(docs_mmr_chinese[0].page_content[:100])
    print()
    print("docs[1]: ")
    print(docs_mmr_chinese[1].page_content[:100])
